#!/usr/bin/env python

from INFILE import siteid,webdir
import os
from datetime import datetime, timedelta
from itertools import groupby
import re
import numpy as np
from numpy import nan
from scipy.stats.stats import nanmean

cldcov = {'CL' : 0,
          'FW' : 1,
          'SC' : 2,
          'PC' : 3,
          'BK' : 4,
          'OV' : 5}


from optparse import OptionParser

tom = datetime.now() + timedelta(hours=24)
tomstr = tom.strftime('%Y%m%d')


# Get the station ID we want from the command line
parser = OptionParser()
parser.add_option('-s', '--site', dest='siteid', action='store', type='string', default=siteid, help='4-letter id of desired site')
parser.add_option('-f', '--fdate', dest='fcst_date',action='store', type='string', default=tomstr, help='Forecast date (YYYYMMDD)')
parser.add_option('-e', '--export', dest='export', action='store_true',default=False, help='Write images to png or not.')
(opts,args) = parser.parse_args()
siteid = opts.siteid.upper()
export_flag = opts.export
fcst_date = datetime.strptime(opts.fcst_date, '%Y%m%d')

if export_flag:
    import matplotlib
    matplotlib.use('agg')
    import matplotlib.pyplot as plt


else:
    import matplotlib
    import matplotlib.pyplot as plt



import pickle


pcode = {'High Temperature' : 'highs',
         'Low Temperature'  : 'lows',
         'QPF code'         : 'qpfc',
         'PoP'              : 'pop',
         'High BC'          : 'hibc',
         'Low BC'           : 'lobc'}

def main():
    parsed = get_data(siteid,fcst_date)
    from plot_current_cmce import get_cmce
    cmce_parsed = get_cmce(siteid, fcst_date)

    plot_ensemble(parsed,cmce_parsed,export_flag)


def bias_correct(memname,hitemp,lotemp,siteid):
    error_dict = pickle.load(open('%s_gfse_stats.pickle' % siteid, 'r'))
    # Find the member number
    if memname.startswith('Control'):
        memnum = 0
    elif memname.strip() == 'Operational':
        return hitemp, lotemp
    else:
        memnum = int(re.search('(\d{1,2})',memname).groups()[0])
   
    # Look up that member number in the error dictionry
    hi_x0 = error_dict[memnum]['hibias_x0']
    hi_x1 = error_dict[memnum]['hibias_x1']
    lo_x0 = error_dict[memnum]['lobias_x0']
    lo_x1 = error_dict[memnum]['lobias_x1']
 
    # Make the correction
    high_bc = hi_x0 + hi_x1 * hitemp
    low_bc = lo_x0 + lo_x1 * lotemp

    return high_bc, low_bc


def get_data(siteid,fcst_date):
    # Download the file from the web
    os.system('wget -O gfse_mos.txt http://www.nws.noaa.gov/cgi-bin/mos/getens.pl?sta=%s' % siteid.upper())

    ensemble_dict = {}
    mos_dict = {}
    for k in pcode.keys():
        ensemble_dict[k] = {}
    
    # Open the file and parse it
    infile = open('gfse_mos.txt','r').read()
    infile_blocks = infile.split('<PRE>')
    for block in infile_blocks:
        block.strip()
        bklines = block.splitlines()
        bkname = bklines[1]
        if not bkname.startswith('Perturbation') and not bkname.startswith('Control') and not bkname.startswith('Oper'):
            continue
     
        # Get the number of the perturbation
        if bkname.startswith('Control'):
            pertnum = 0
        elif bkname.startswith('Perturbation'):
            pertnum = int(re.search('(\d{1,2})',bkname).groups()[0])
        else:
            pertnum = 999

        # Get the forecast date
        analysis_date = re.search('(\d{1,2}/\d{1,2}/\d{4})',bklines[2]).groups()[0]
        analysis_dt = datetime.strptime(analysis_date,'%m/%d/%Y')
        print "DATE FOUND:", analysis_dt
        dates = bklines[3].split('|')
        forecast_dts = []
        for date in dates:
            datesp = date.split()
            fcst_time = analysis_dt + timedelta(hours=int(datesp[-1])-24)
            forecast_dts.append(fcst_time)       
        #print forecast_dts
        

        # Now we have a list of dates.  Use this to key the different variables
        # Start with temperatures
        xn = bklines[5].split('|')
        xn[0] = xn[0][4:]
        num = xn[forecast_dts.index(fcst_date)]
        if len(num.split()) == 2:
            low, high = num.split()
            low = int(low)
            high = int(high)
        elif len(num.split()) == 1:
            high = int(num)
            low = nan
        else:
            high = nan
            low = nan

        

        ensemble_dict['High Temperature'][bkname] = high
        ensemble_dict['Low Temperature'][bkname] = low
        mos_dict['m%d_high' % pertnum] = high
        mos_dict['m%d_low' % pertnum] = low
        high_bc, low_bc = bias_correct(bkname,high,low,siteid)
        ensemble_dict['High BC'][bkname] = high_bc
        ensemble_dict['Low BC'][bkname] = low_bc
 


        # Now for precip
        # Have to check for operational (it's on a different line)
        if bkname == 'Operational':
            ddex = 7
            cdex = 8
            pdex = 11
            qdex = 13
        else:
            ddex = 7
            cdex = 8
            pdex = 10
            qdex = 12
        p24 = bklines[pdex].split('|')
        num = p24[forecast_dts.index(fcst_date)]
        try:
            p24 = int(num)
        except:
            p24 = nan

        # And qpf
        q24 = bklines[qdex].split('|')
        num = q24[forecast_dts.index(fcst_date)]
        try:
            qpf24 = int(num)
        except:
            qpf24 = nan

        ensemble_dict['PoP'][bkname] = p24
        ensemble_dict['QPF code'][bkname] = qpf24
        mos_dict['m%d_p24' % pertnum] = p24
        mos_dict['m%d_q24' % pertnum] = qpf24

        # for the MOS computation, grab dewpoint and cloud as well
        dpt = bklines[ddex].split('|')
        dpt[0] = dpt[0][4:]
        num = dpt[forecast_dts.index(fcst_date)]
        if len(num.split()) == 2:
            dpt1, dpt2 = num.split()
            dpt1 = int(dpt1)
            dpt2 = int(dpt2)
        elif len(num.split()) == 1:
            dpt2 = int(num)
            dpt1 = nan
        else:
            dpt2 = nan
            dpt1 = nan
        dptval = nanmean([dpt1,dpt2]) 
        mos_dict['m%d_dpt' % pertnum] = dptval


        cld = bklines[cdex].split('|')
        cld[0] = cld[0][4:]
        num = cld[forecast_dts.index(fcst_date)]
        if len(num.split()) == 2:
            cld1, cld2 = num.split()
            cld1 = cldcov[cld1]
            cld2 = cldcov[cld2]
        elif len(num.split()) == 1:
            cld2 = cldcov[num]
            cld1 = nan
        else:
            cld2 = nan
            cld1 = nan
        cldval = nanmean([cld1,cld2])
        mos_dict['m%d_cld' % pertnum] = cldval


	print bkname, fcst_date, high, low, p24, qpf24


    from gfse_mos_analysis import evaluate_fcst
    mos_vdict = evaluate_fcst(mos_dict,fcst_date)
    #raw_input()
    return ensemble_dict



def BMA_compute(ensemble_dict,curvex,varname,siteid):
    # Go through each ensemble member and compute the
    # weighted BMA curve of each

    error_dict = pickle.load(open('%s_gfse_stats.pickle' % siteid, 'r'))
    outline = np.zeros(np.shape(curvex))
    out_fcst = 0.
    for memname in ensemble_dict['High BC'].keys():
        # Find the member number
        if memname.startswith('Control'):
            memnum = 0
        elif memname.strip() == 'Operational':
            continue
        else:
            memnum = int(re.search('(\d{1,2})',memname).groups()[0])

        if varname.upper().startswith('HIGH'):
            std = error_dict[memnum]['hierr']['spread']
            weight = error_dict[memnum]['closest_hi']
            meanval = ensemble_dict['High BC'][memname]
        elif varname.upper().startswith('LOW'):
            std = error_dict[memnum]['loerr']['spread']
            weight = error_dict[memnum]['closest_lo']
            meanval = ensemble_dict['Low BC'][memname]
        else:
            print "Unrecognized BMA varname:", varname
            exit(1) 

        # Now compute the curve
        norm_vals = 1/(std * np.sqrt(2 * np.pi)) * np.exp(-(curvex-meanval)**2 / (2*std**2))
        # And add it to the BMA
        outline = outline + norm_vals * weight
        out_fcst = out_fcst + weight * meanval
    return outline, out_fcst

def gamma_dist_compute(mean,std,xvals):
    # Compute a two-parameter gamma distribution over xvals
    from scipy.special import gamma
    beta = mean/(std**2)
    alpha = beta * mean

    # Now compute the curve
    curve_vals = [(beta**alpha / gamma(alpha)) * x ** (alpha-1) * np.exp(-beta*x) for x in xvals]
    return curve_vals


def plot_ensemble(ensemble_dict, cmce_dict, export_flag):
    # Use the ensemble dictionary to make plots
    # Need to compute normal distribution

    for var in ensemble_dict.keys():
        if var.endswith('BC'):
            continue
        varlist = [ensemble_dict[var][k] for k in ensemble_dict[var].keys()]
        if var.endswith('Temperature'):
            cmce_varlist = [cmce_dict[var][k] for k in cmce_dict[var].keys()]
            median_cmce = np.median(cmce_varlist)

        mean_var = np.mean(varlist)
        median_var = np.median(varlist)
        std_var = np.std(varlist)

        # Compute the BMA curve
        curve_xvals = np.linspace(min(varlist)-7, max(varlist)+7,100)
        if var.startswith('High') or var.startswith('Low'):
            BMA_vals, BMA_fcst = BMA_compute(ensemble_dict,curve_xvals,var,siteid)  
        else:
            gamma_vals = gamma_dist_compute(mean_var,std_var,curve_xvals)
        
        # Find the operational value
        op_value = ensemble_dict[var]['Operational']


        # Make a dummy plot just to get the cmce hist values
        if var.endswith('Temperature'):
            binrge = range(int(min(varlist + cmce_varlist)),int(max(varlist + cmce_varlist))+2)
            nbinrge = [r-0.5 for r in binrge]
            #binrge = nbinrge[:]

            n_cmce, bins_cmce = np.histogram(cmce_varlist,bins=nbinrge)
            n_gfse, bins_gfse = np.histogram(varlist,bins=nbinrge)
        else:
            binrge = range(int(min(varlist)),int(max(varlist))+2)
            nbinrge = [r-0.5 for r in binrge]
            binrge = nbinrge[:]
            n_gfse, bins_gfse = np.histogram(varlist,bins=binrge)

        # plot everything
        plt.figure(figsize=(12,11))
        ax1 = plt.gca()
        if var.endswith('Temperature'):
             patches = ax1.bar(binrge[:-1],n_gfse,align='center',color='LightSkyBlue',label='GFSE Hist.')
             ax1.hold(True)
             patches_cmc = ax1.bar(binrge[:-1],n_cmce,align='center',bottom=n_gfse,color='LightGreen',label='CMCE Hist.')  
        else:
             n,bin,patches = ax1.hist(varlist,bins=binrge,color='LightSkyBlue',label='GFSE Hist.')

        plt.ylabel('Number of ensemble members')
        plt.xlabel(var)
        ax2 = ax1.twinx()
        #plt.scatter(varlist,np.zeros(len(varlist)))
        median = ax2.axvline(x=median_var,c='b', linewidth=3, label='GFSE Median')
        #mean = ax2.axvline(x=mean_var,c='b', linestyle='dashed', linewidth=3, label='Mean')
        oper = ax2.axvline(x=op_value,c='k',linewidth=4, label='Oper. GFS MOS')
        if var.startswith('High') or var.startswith('Low'):
            cmce_median = ax2.axvline(x=median_cmce,c='g', linewidth=3, label='CMCE Median')
            bmamax = max(BMA_vals)
            left_90, right_90, fill_xvals_90, fill_bma_vals_90 = find_percentiles(curve_xvals, BMA_vals, 90.)
            left_50, right_50, fill_xvals_50, fill_bma_vals_50 = find_percentiles(curve_xvals, BMA_vals, 50.)


            bma = ax2.plot(curve_xvals, BMA_vals, c='m',linestyle='dotted', linewidth=3,label= 'BMA Curve')
            bma_90 = ax2.fill_between(fill_xvals_90,y1=fill_bma_vals_90,y2=0, color='m',alpha=0.1)
            bma_50 = ax2.fill_between(fill_xvals_50,y1=fill_bma_vals_50,y2=0, color='m',alpha=0.3)
            bma_horiz = ax2.axhline(y=bmamax,c='m',linewidth=1,label='BMA Max. Prob.')
            bma_fcst = ax2.axvline(x=BMA_fcst,c='m',linewidth=3, label='BMA Fcst')
            ax2.set_ylim((0,0.25)) 
            leg = plt.legend([bma[0],bma_fcst,median,cmce_median,oper,patches[0],patches_cmc[0]],\
                       ['BMA curve','BMA fcst.','GFSE Median','CMCE Median','Oper. GFS MOS','GFSE Hist.','CMCE Hist.'],\
                        loc=0,ncol=4,mode='expand')
            #leg = pylab.legend(modlist+moslist,loc=8,ncol=5,mode='expand')


            # Make the legend transparent
            leg.get_frame().set_alpha(0.5)
            ltexts = leg.get_texts()
            plt.setp(ltexts,fontsize='small')
 

            ax = plt.gca()
            ax.xaxis.set_major_locator(matplotlib.ticker.MultipleLocator(1.0))
            # Do some exporting here
            if os.path.exists('%s_gfse_performance.pickle' % siteid.upper()):
                gfse_perf = pickle.load(open('%s_gfse_performance.pickle' % siteid.upper(), 'r'))
            else:
                gfse_perf = {}
            if fcst_date not in gfse_perf.keys():
                gfse_perf[fcst_date] = {}
            gfse_perf[fcst_date][pcode[var] + '_oper'] = op_value
            gfse_perf[fcst_date][pcode[var] + '_median'] = median_var
            gfse_perf[fcst_date][pcode[var] + '_bma'] = BMA_fcst
            gfse_perf[fcst_date][pcode[var] + '_low50'] = left_50
            gfse_perf[fcst_date][pcode[var] + '_high50'] = right_50
            gfse_perf[fcst_date][pcode[var] + '_low90'] = left_90
            gfse_perf[fcst_date][pcode[var] + '_high90'] = right_90
            gfse_perf[fcst_date][pcode[var] + '_gfsemems'] = varlist
            gfse_perf[fcst_date][pcode[var] + '_cmcemems'] = cmce_varlist
            pickle.dump(gfse_perf,open('%s_gfse_performance.pickle' % siteid.upper(),'w'))
            plt.title('NAEFS Ensemble Forecast of %s at %s for %s' % (var,siteid.upper(),fcst_date.strftime('%m/%d/%Y')))
           


        else:
            gammav = ax2.plot(curve_xvals, gamma_vals, c='g', linestyle='dotted', linewidth=3,label= 'Norm. Curve')
            left_90, right_90, fill_xvals_90, fill_gamma_vals_90 = find_percentiles(curve_xvals,gamma_vals,90.)
            left_50, right_50, fill_xvals_50, fill_gamma_vals_50 = find_percentiles(curve_xvals,gamma_vals,50.)
            gam_90 = ax2.fill_between(fill_xvals_90, y1=fill_gamma_vals_90, y2=0, color='g', alpha=0.1)
            gam_50 = ax2.fill_between(fill_xvals_50, y1=fill_gamma_vals_50, y2=0, color='g', alpha=0.3)
            plt.legend([gammav[0],median,oper,patches[0]],['Gamma Curve','Median','Oper. GFS MOS','Ens. Hist.'],loc=0)
            ax2.set_ylim((0,0.7)) 
            if var.startswith('QPF'):
                # Change xticks
                plt.xticks([0,1,2,3,4,5],['0.00 in.','0.01-0.09 in.','0.10-0.24 in.','0.25-0.49 in.','0.50-0.99 in.','>1.00 in.'])
                ax2.set_xlim((0,6))
            else:
                # For PoP, only set left limit of zero if the minimum actually is zero
                if min(varlist) == 0:
                    x1,x2,y1,y2 = plt.axis()
                    ax2.set_xlim((0.0,x2))
            #plt.scatter(varlist,np.zeros(len(varlist)))
            plt.title('GFS Ensemble Forecast of %s at %s for %s' % (var,siteid.upper(),fcst_date.strftime('%m/%d/%Y')))
        if export_flag:
            plt.savefig('%s_GFSE_%s.png' % (siteid.upper(),pcode[var]), bbox_inches='tight')
            os.system('mv %s_GFSE_%s.png %s' % (siteid.upper(), pcode[var], webdir)) 
        else:
            plt.show()

def find_percentiles(xvals,yvals,pctile):
    dp = [r * (xvals[-1] - xvals[-2]) for r in yvals]
    accum_sum = [np.nansum(dp[0:r]) for r in range(len(dp))]
    #print accum_sum
    #print "SUM:", accum_sum[-1]

    left_found = False
    right_found = False
    left_val = nan
    right_val = nan
   
    fill_xvals = []
    fill_yvals = []
    low_pct = (100.-pctile)/(200.)
    up_pct = 1.0 - low_pct
    for r,num,BMAv in zip(accum_sum, xvals, yvals):
        if r >= low_pct and not left_found:
            left_val = num
            fill_xvals.append(num)
            fill_yvals.append(BMAv)
            left_found = True
        if r > up_pct and not right_found:
            right_val = num
            right_found = True

        if left_found and not right_found:
            fill_xvals.append(num)
            fill_yvals.append(BMAv)

    return left_val, right_val, fill_xvals, fill_yvals



if __name__ == '__main__':
    main()
